# Copyright © 2022 BAAI. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License")

import os
import json
import time
import shlex
import socket
import argparse
import subprocess
from collections import OrderedDict


def run_local_cmd(command, check=True, preview=False, blocking=True, verbose=True):
    ''' Runs command via subprocess in the local machine.

    Args:
        command: command to run on the host
        check: (optional) whether to check the return code of the command
        verbose: (optional) whether to print the command

    Returns:
        result of the command, return code
    '''

    if verbose:
        print("Running command: {} on localhost".format(command), flush=True)

    if not preview:
        if not blocking:
            process = subprocess.Popen(command,
                                       shell=True,
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)
            return process, command 
        else:
            result = subprocess.run(command,
                                    shell=True,
                                    capture_output=True,
                                    check=check)
            if verbose:
                if result.returncode != 0:
                    print("Command {} failed on localhost".format(command), flush=True)
                else:
                    print("Command {} successfully executed on localhost".format(command), flush=True)
            return result.stdout.decode("utf-8"), result.returncode
    else:
        return None, None


def run_remote_cmd(host, command, user=None, password=None, blocking=True, preview=False, verbose=True):
    ''' Runs ssh command via subprocess.

    Args:
        host: target host to send the command to
        command: command to run on the host
        user: (optional) user to use to login to host
        password: (optional) password for ssh login
        check: (optional) whether to check the return code of the command
        verbose: (optional) whether to print the command

    Returns:
        result of the command, return code
    '''

    where = "%s" % host if user is None else "%s@%s" % (user, host)
    if password is not None:
        prefix = "sshpass -p {} ".format(password)
    else:
        prefix = ""
    ssh_cmd = "ssh {} {}".format(where, command)
    if verbose:
        print("Running command: {} on host {}".format(
            ssh_cmd, host), flush=True)

    if not preview:
        if not blocking:
            process = subprocess.Popen(prefix + ssh_cmd,
                                      shell=True,
                                      stdout=subprocess.PIPE,
                                      stderr=subprocess.PIPE)
            return process, ssh_cmd
        else:
            result = subprocess.run(prefix + ssh_cmd,
                                    shell=True,
                                    capture_output=True,
                                    check=False)
            if verbose:
                if result.returncode != 0:
                    print("Command {} failed on host {}".format(ssh_cmd, host), flush=True)
                else:
                    print("Command {} successfully executed on host {}".format(ssh_cmd,host), flush=True)
            return result.stdout.decode("utf-8"), result.returncode
    else:
        return None, None


def run_scp_cmd(host, src, dst, password=None, user=None, blocking=True, preview=False, verbose=True):
    ''' Runs scp command via subprocess.

    Args:
        host: target host to send the command to
        src: source file or directory
        dst: destination file or directory
        password: (optional) password for ssh login
        user: (optional) user to use to login to host
        check: (optional) whether to check the return code of the command
        verbose: (optional) whether to print the command

    Returns:
        result of the command, return code
    '''

    where = "%s" % host if user is None else "%s@%s" % (user, host)
    if password is not None:
        prefix = "sshpass -p {} ".format(password)
    else:
        prefix = ""
    scp_cmd = "scp -r {}:{} {}".format(where, src, dst)
    if verbose:
        print("Running {}".format(scp_cmd), flush=True)

    if not preview:
        if not blocking:
            process = subprocess.Popen(prefix + scp_cmd,
                                      shell=True,
                                      stdout=subprocess.PIPE,
                                      stderr=subprocess.PIPE)
            return process, scp_cmd 
        else:
            result = subprocess.run(prefix + scp_cmd,
                                    shell=True,
                                    capture_output=True,
                                    check=True)
            if verbose:
                if result.returncode != 0:
                    print("Command {} failed on host {}".format(scp_cmd, host), flush=True)
                else:
                    print("Command {} successfully executed on host {}".format(scp_cmd,host), flush=True)
            return result.stdout.decode("utf-8"), result.returncode
    else:
        return None, None 


def get_ip():
    """Get the IP address of the current host."""
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.settimeout(0)
    try:
        # doesn't even have to be reachable
        s.connect(('10.254.254.254', 1))
        IP = s.getsockname()[0]
    except Exception:
        IP = '127.0.0.1'
    finally:
        s.close()
    return IP


def get_args():
    """Parse commandline."""

    parser = argparse.ArgumentParser(description='Pull remote checkpoints')
    parser.add_argument('--ckpt-dir', type=str, required=True,
                        help='Directory to load model checkpoint from')
    parser.add_argument('--saved-dir', type=str, default=None,
                        help='Directory to save model checkpoint to')
    parser.add_argument('--iter', type=int, required=True,
                        help='Iteration of the checkpoints to use')
    parser.add_argument('--hostfile', type=str, default=None,
                        help='Path to the hostfile')
    parser.add_argument('--password', type=str, default=None,
                        help='password for ssh login')
    parser.add_argument('--all', action='store_true',
                        help='Pull all checkpoints to current host')
    parser.add_argument('--type', default='all',
                        choices=['all', 'required'],
                        help='The amount of checkpoints to be pulled.')
    parser.add_argument('--preview', action='store_true',
                        help='Preview the checkpoints to be pulled')
    args = parser.parse_args()

    return args


def pull_required_ckpts(args):
    ckpt_dir = args.ckpt_dir
    saved_dir = args.saved_dir
    iter = args.iter
    password = args.password
    cur_host = get_ip()

    iter_dir = os.path.join(ckpt_dir, 'iter_{:07d}'.format(iter))
    if saved_dir and ckpt_dir != saved_dir:
        saved_iter_dir = os.path.join(saved_dir, 'iter_{:07d}'.format(iter))
    else:
        saved_dir = ckpt_dir
        saved_iter_dir = iter_dir
    if not os.path.exists(saved_iter_dir):
        os.makedirs(saved_iter_dir)

    ckpts_info_path = os.path.join(ckpt_dir, "ckpt_info.json")
    assert os.path.exists(
        ckpts_info_path), "ckpt_info.json does not exist in {}".format(ckpt_dir)
    with open(ckpts_info_path, 'r') as f:
        ckpts_info = json.load(f)

    ckpt_to_host = {}
    host_to_ranks = {}
    rank_to_ckpt = {}
    for ckpt_info in ckpts_info:
        if ckpt_info['existed']:
            if ckpt_info['ckpt_dir'] in ckpt_to_host \
                    and ckpt_info['host'] != ckpt_to_host[ckpt_info['ckpt_dir']]:
                print("Warning: checkpoint {} exists on multiple hosts".format(
                    ckpt_info['ckpt_dir']), flush=True)
            ckpt_to_host[ckpt_info['ckpt_dir']] = ckpt_info['host']

        if host_to_ranks.get(ckpt_info['host'], None) is None:
            host_to_ranks[ckpt_info['host']] = []
        host_to_ranks[ckpt_info['host']].append(ckpt_info['rank'])

        rank_to_ckpt[ckpt_info['rank']] = ckpt_info['ckpt_dir']

    required_ckpts = set()
    for rank in host_to_ranks[cur_host]:
        required_ckpts.add(rank_to_ckpt[rank])

    # Pull the checkpoints from the hosts
    processes = []
    cmds = []
    for ckpt in required_ckpts:
        ckpt_host = ckpt_to_host[ckpt]
        src = os.path.join(iter_dir, ckpt)
        dst = saved_iter_dir
        if cur_host != ckpt_host:
            process, cmd = run_scp_cmd(ckpt_host, src, dst, password,
                                  preview=args.preview, blocking=False)
            processes.append(process)
            cmds.append(cmd)
        else:
            if os.path.dirname(src) != dst:
                process, cmd = run_local_cmd("cp -r {} {}".format(src, dst), preview=args.preview, blocking=False)
                processes.append(process)
                cmds.append(cmd)
            else:
                print("Warning: checkpoint {} already exists on {}.".format(
                    src, cur_host), flush=True)
    visited = {process: False for process in processes} 
    while not all(visited.values()):
        for i, process in enumerate(processes):
            if process is None:
                visited[process] = True
                continue
            if process.poll() is not None and visited[process] is False:
                visited[process] = True
                _, error = process.communicate()
                if error:
                    print("Error {} when executing command {} on host {}".format(error, cmds[i], cur_host), flush=True)
        time.sleep(5)

    tracker_file = os.path.join(saved_dir, 'latest_checkpointed_iteration.txt')
    with open(tracker_file, 'w') as f:
        f.write(str(iter))

    print("Finished pulling the required checkpoints", flush=True)


def pull_all_ckpts(args):
    ckpt_dir = args.ckpt_dir
    saved_dir = args.saved_dir
    iter = args.iter
    hostfile = args.hostfile
    password = args.password
    cur_host = get_ip()

    assert hostfile is not None, "hostfile must be provided when pulling all checkpoints"
    iter_dir = os.path.join(ckpt_dir, 'iter_{:07d}'.format(iter))
    if saved_dir and ckpt_dir != saved_dir:
        saved_iter_dir = os.path.join(saved_dir, 'iter_{:07d}'.format(iter))
    else:
        saved_dir = ckpt_dir
        saved_iter_dir = iter_dir
    if not os.path.exists(saved_iter_dir):
        os.makedirs(saved_iter_dir)

    # Parse the host file and get host list
    hostlist = []
    with open(hostfile, 'r') as f:
        for line in f:
            host_and_slot = line.strip().split()
            if len(host_and_slot) > 0:
                host = line.strip().split()[0]
                hostlist.append(host)

    # Build a dictionary of host to checkpoints
    host_to_ckpts = OrderedDict()
    for host in hostlist:
        if host not in host_to_ckpts:
            host_to_ckpts[host] = []
        result = run_remote_cmd(
            host, "test -d " + iter_dir, password=password)
        if result[1] != 0:
            print("Warning: directory or file {} does not exists on host {}".format(
                "", host), flush=True)
            continue
        else:
            result = run_remote_cmd(
                host, "ls {}".format(iter_dir), password=password)
            for line in result[0].splitlines():
                inner_dir = os.path.join(iter_dir, line)
                inner_result = run_remote_cmd(
                    host, "ls {}".format(inner_dir), password=password)
                if len(inner_result[0]) != 0:
                    host_to_ckpts[host].append(line)

    # Pull the checkpoints from the hosts
    processes = []
    cmds = []
    visited_ckpts = set()
    for host, ckpts in host_to_ckpts.items():
        for ckpt in ckpts:
            if ckpt not in visited_ckpts:
                visited_ckpts.add(ckpt)
                src = os.path.join(iter_dir, ckpt)
                dst = os.path.join(saved_iter_dir, ckpt)
                if cur_host != host:
                    process, cmd = run_scp_cmd(host, src, dst, password, preview=args.preview, blocking=False)
                    processes.append(process)
                    cmds.append(cmd)
                else:
                    if src != dst:
                        process, cmd = run_local_cmd("cp {} {}".format(os.path.join(iter_dir, ckpt), os.path.join(
                            saved_iter_dir, ckpt)), preview=args.preview, blocking=False)
                        processes.append(process)
                        cmds.append(cmd)
            else:
                print("Warning: checkpoint {} on {} also exists on other hosts.".format(
                    ckpt, host), flush=True)

    visited = {process: False for process in processes} 
    while not all(visited.values()):
        for i, process in enumerate(processes):
            if process is None:
                visited[process] = True
                continue
            if process.poll() is not None and visited[process] is False:
                visited[process] = True
                _, error = process.communicate()
                if error:
                    print("Error {} when executing command {} on host {}".format(error, cmds[i], cur_host), flush=True)
        time.sleep(5)
    tracker_file = os.path.join(saved_dir, 'latest_checkpointed_iteration.txt')
    with open(tracker_file, 'w') as f:
        f.write(str(iter))

    print("Finished pulling all checkpoints.", flush=True)


if __name__ == '__main__':
    args = get_args()

    if args.type == "all":
        pull_all_ckpts(args)
    elif args.type == "required":
        pull_required_ckpts(args)
    else:
        raise NotImplementedError
